#include "Render/PiccoloRenderTarget.h"

#include "ECS/Component/PLookAtCameraComponent.h"

#include "PiccoloApplication.h"
#include "Graphics/PiccoloVKRenderPass.h"
#include "Graphics/PiccoloVKImage.h"

namespace Piccolo
{
    PiccoloRenderTarget::PiccoloRenderTarget(PiccoloVKRenderPass *renderPass)
    {
        PiccoloRenderContext* renderContext = PiccoloApplication::GetAppContext()->renderContext;
        PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();

        mRenderPass = renderPass;
        mExtent = { swapchain->GetWidth(), swapchain->GetHeight() };
        mBufferCount = swapchain->GetImages().size();
        bSwapchainTarget = true;

        Init();
        Recreate();
    }

    PiccoloRenderTarget::PiccoloRenderTarget(PiccoloVKRenderPass *renderPass, uint32_t bufferCount, VkExtent2D extent)
        : mRenderPass(renderPass), mBufferCount(bufferCount), mExtent(extent), bSwapchainTarget(false)
    {
        Init();
        Recreate();
    }

    PiccoloRenderTarget::~PiccoloRenderTarget()
    {
        for (const auto& item : mMaterialSystemList)
        {
            item->OnDestroy();
        }
        mMaterialSystemList.clear();
    }

    void PiccoloRenderTarget::Init()
    {
        mClearValues.resize(mRenderPass->GetAttachmentSize());
        SetColorClearValue({ 0.f, 0.f, 0.f, 1.f });
        SetDepthStencilClearValue({ 1.f, 0 });
    }

    void PiccoloRenderTarget::Recreate()
    {
        if (mExtent.width == 0 || mExtent.height == 0) return;
        mFrameBuffers.clear();
        mFrameBuffers.resize(mBufferCount);

        PiccoloRenderContext* renderContext = PiccoloApplication::GetAppContext()->renderContext;
        PiccoloVKDevice* device = renderContext->GetDevice();
        PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();

        std::vector<Attachment> attachments = mRenderPass->GetAttachments();
        if (attachments.empty()) return;

        std::vector<VkImage> swapchainImages = swapchain->GetImages();

        for (int i = 0; i < mBufferCount; ++i)
        {
            std::vector<std::shared_ptr<PiccoloVKImage>> images;
            for (int j = 0; j < attachments.size(); ++j)
            {
                Attachment attachment = attachments[j];
                if (bSwapchainTarget && attachment.finalLayout == VK_IMAGE_LAYOUT_PRESENT_SRC_KHR && attachment.samples == VK_SAMPLE_COUNT_1_BIT)
                {
                    images.push_back(std::make_shared<PiccoloVKImage>(device, swapchainImages[i], VkExtent3D{ mExtent.width, mExtent.height, 1 }, attachment.format, attachment.usage));
                }
                else
                {
                    images.push_back(std::make_shared<PiccoloVKImage>(device, VkExtent3D{ mExtent.width, mExtent.height, 1 }, attachment.format, attachment.usage, attachment.samples));
                }
            }

            mFrameBuffers[i] = std::make_shared<PiccoloVKFrameBuffer>(device, mRenderPass, images, mExtent.width, mExtent.height);
            images.clear();
        }
    }

    void PiccoloRenderTarget::Begin(VkCommandBuffer cmdBuffer)
    {
        assert(!bBeginTarget && "You should not called Begin() again.");
        if (bShouldUpdate)
        {
            Recreate();
            bShouldUpdate = false;
        }

        if (PEntity::HasComponent<PLookAtCameraComponent>(mCamera))
        {
            mCamera->GetComponent<PLookAtCameraComponent>().SetAspect(mExtent.width * 1.f / mExtent.height);
        }

        if (bSwapchainTarget)
        {
            PiccoloRenderContext* renderContext = PiccoloApplication::GetAppContext()->renderContext;
            PiccoloVKSwapchain* swapchain = renderContext->GetSwapchain();
            mCurrentBufferIndex = swapchain->GetCurrentImageIndex();
        }
        else
        {
            mCurrentBufferIndex = (mCurrentBufferIndex + 1) % mBufferCount;
        }

        mRenderPass->BeginRenderPass(cmdBuffer, GetCurrentFrameBuffer(), mClearValues);
        bBeginTarget = true;
    }

    void PiccoloRenderTarget::End(VkCommandBuffer cmdBuffer)
    {
        if (bBeginTarget)
        {
            mRenderPass->EndRenderPass(cmdBuffer);
            bBeginTarget = false;
        }
    }

    void PiccoloRenderTarget::SetExtent(const VkExtent2D &extent)
    {
        mExtent = extent;
        bShouldUpdate = true;
    }

    void PiccoloRenderTarget::SetBufferCount(uint32_t bufferCount)
    {
        mBufferCount = bufferCount;
        bShouldUpdate = true;
    }

    void PiccoloRenderTarget::SetColorClearValue(VkClearColorValue colorClearValue)
    {
        std::vector<Attachment> renderPassAttachments = mRenderPass->GetAttachments();
        for (int i = 0; i < renderPassAttachments.size(); ++i)
        {
            if (!IsDepthStencilFormat(renderPassAttachments[i].format) && renderPassAttachments[i].loadOp == VK_ATTACHMENT_LOAD_OP_CLEAR)
            {
                mClearValues[i].color = colorClearValue;
            }
        }
    }

    void PiccoloRenderTarget::SetColorClearValue(uint32_t attachmentIndex, VkClearColorValue colorClearValue)
    {
        std::vector<Attachment> renderPassAttachments = mRenderPass->GetAttachments();
        if (attachmentIndex <= renderPassAttachments.size() - 1)
        {
            if (!IsDepthStencilFormat(renderPassAttachments[attachmentIndex].format) && renderPassAttachments[attachmentIndex].loadOp == VK_ATTACHMENT_LOAD_OP_CLEAR)
            {
                mClearValues[attachmentIndex].color = colorClearValue;
            }
        }
    }

    void PiccoloRenderTarget::SetDepthStencilClearValue(VkClearDepthStencilValue depthStencilValue)
    {
        std::vector<Attachment> renderPassAttachments = mRenderPass->GetAttachments();
        for (int i = 0; i < renderPassAttachments.size(); ++i)
        {
            if (IsDepthStencilFormat(renderPassAttachments[i].format) && renderPassAttachments[i].loadOp == VK_ATTACHMENT_LOAD_OP_CLEAR)
            {
                mClearValues[i].depthStencil = depthStencilValue;
            }
        }
    }

    void PiccoloRenderTarget::SetDepthStencilClearValue(uint32_t attachmentIndex, VkClearDepthStencilValue depthStencilValue)
    {
        std::vector<Attachment> renderPassAttachments = mRenderPass->GetAttachments();
        if (attachmentIndex <= renderPassAttachments.size() - 1)
        {
            if (IsDepthStencilFormat(renderPassAttachments[attachmentIndex].format) && renderPassAttachments[attachmentIndex].loadOp == VK_ATTACHMENT_LOAD_OP_CLEAR)
            {
                mClearValues[attachmentIndex].depthStencil = depthStencilValue;
            }
        }
    }
}
